import rclpy
from rclpy.node import Node
from rclpy.time import Time
from geometry_msgs.msg import PoseWithCovarianceStamped
import numpy as np

class SensorFusionROS2(Node):
    def __init__(self):
        super().__init__('sensor_fusion_node')
        
        # 创建订阅者
        self.gps_sub = self.create_subscription(
            PoseWithCovarianceStamped,
            'gps_pose',
            self.gps_callback,
            10
        )
        
        self.imu_sub = self.create_subscription(
            PoseWithCovarianceStamped,
            'imu_pose',
            self.imu_callback,
            10
        )
        
        self.odom_sub = self.create_subscription(
            PoseWithCovarianceStamped,
            'odom_pose',
            self.odom_callback,
            10
        )
        
        # 创建发布者
        self.fused_pose_pub = self.create_publisher(
            PoseWithCovarianceStamped,
            'fused_pose',
            10
        )
        
        # 回调函数
        self.gps_callback_fn = None
        self.imu_callback_fn = None
        self.odom_callback_fn = None
        self.state_callback_fn = None
        
    
    def set_gps_callback(self, callback):
        """设置GPS回调函数"""
        self.gps_callback_fn = callback
    
    def set_imu_callback(self, callback):
        """设置IMU回调函数"""
        self.imu_callback_fn = callback
    
    def set_odom_callback(self, callback):
        """设置里程计回调函数"""
        self.odom_callback_fn = callback
    
    def set_state_callback(self, callback):
        """设置状态回调函数"""
        self.state_callback_fn = callback
    
    def gps_callback(self, msg):
        """GPS消息回调"""
        if self.gps_callback_fn:
            timestamp = Time.from_msg(msg.header.stamp).nanoseconds * 1e-9
            self.gps_callback_fn(msg, timestamp)
    
    def imu_callback(self, msg):
        """IMU消息回调"""
        if self.imu_callback_fn:
            timestamp = Time.from_msg(msg.header.stamp).nanoseconds * 1e-9
            self.imu_callback_fn(msg, timestamp)
    
    def odom_callback(self, msg):
        """里程计消息回调"""
        if self.odom_callback_fn:
            timestamp = Time.from_msg(msg.header.stamp).nanoseconds * 1e-9
            self.odom_callback_fn(msg, timestamp)
    
    def publish_fused_pose(self, state, covariance=None):
        """发布融合后的位姿"""
        if state is None:
            return
            
        # 创建消息
        msg = PoseWithCovarianceStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'map'  # 假设使用map坐标系
        
        # 设置位置
        msg.pose.pose.position.x = state[0, 0]
        msg.pose.pose.position.y = state[1, 0]
        msg.pose.pose.position.z = 0.0
        
        # 设置朝向 (四元数)
        theta = state[2, 0]
        msg.pose.pose.orientation.x = 0.0
        msg.pose.pose.orientation.y = 0.0
        msg.pose.pose.orientation.z = np.sin(theta / 2)
        msg.pose.pose.orientation.w = np.cos(theta / 2)
        
        # 设置协方差 (简化版)
        if covariance is not None:
            # 提取位置和朝向的协方差
            pose_covariance = [0.0] * 36
            for i in range(6):
                for j in range(6):
                    pose_covariance[i*6 + j] = covariance[i, j]
            msg.pose.covariance = pose_covariance
        else:
            # 默认协方差
            msg.pose.covariance = [
                0.1, 0.0, 0.0, 0.0, 0.0, 0.0,
                0.0, 0.1, 0.0, 0.0, 0.0, 0.0,
                0.0, 0.0, 0.1, 0.0, 0.0, 0.0,
                0.0, 0.0, 0.0, 0.1, 0.0, 0.0,
                0.0, 0.0, 0.0, 0.0, 0.1, 0.0,
                0.0, 0.0, 0.0, 0.0, 0.0, 0.1
            ]
        
        # 发布消息
        self.fused_pose_pub.publish(msg)    
